// Copyright (c) 2021, the Dart project authors.  Please see the AUTHORS file
// for details. All rights reserved. Use of this source code is governed by a
// BSD-style license that can be found in the LICENSE file.

import 'package:vm/testing/il_matchers.dart';

bool shouldPrint = false;

@pragma('vm:never-inline')
void blackhole(Object v) {
  if (shouldPrint) {
    print(v);
  }
}

class A {}

class A0 extends A {
  final double x;
  final String str;

  A0(this.x, this.str);

  // Use [x] to prevent the field from being shaken out. We would like
  // [x] to occupy the same location that [A1.str] takes so that
  // type confusion / incorrect LICM would cause a crash when we load
  // `o.[A1.str].length` on an object of type [A0]
  String toString() => 'A0($x)';
}

class A1 extends A {
  final String str;

  A1(this.str);
}

class H<T> {
  final T data;
  H(this.data);
}

abstract class B<T extends A> {
  final T v;

  B(this.v);

  int load(H<T> h);
  int loadWithNamedParam({required H<T> h});

  @pragma('vm:never-inline')
  @pragma('vm:testing:print-flow-graph', '*LICM')
  int testNarrowingThroughThisCallWithPositionalParam(H<T> h) {
    var result = 0;
    for (var i = 0; i < 10; i++) {
      // We will perform polymorphic inlining of `load` because `this`
      // is known to be either B0 or B1. In both cases we will have
      // v.str.length loads fully inlined because inlined bodies
      // have precise type information for v.
      // Then we will hoist v.str.length out of the loop past
      // class-id comparisons generated by the inlining leading
      // to incorrect code which will crash.
      result += load(h);
    }
    return result;
  }

  @pragma('vm:never-inline')
  @pragma('vm:testing:print-flow-graph', '*LICM')
  int testNarrowingThroughThisCallWithNamedParams(H<T> h) {
    var result = 0;
    for (var i = 0; i < 10; i++) {
      // We will perform polymorphic inlining of `load` because `this`
      // is known to be either B0 or B1. In both cases we will have
      // v.str.length loads fully inlined because inlined bodies
      // have precise type information for v.
      // Then we will hoist v.str.length out of the loop past
      // class-id comparisons generated by the inlining leading
      // to incorrect code which will crash.
      result += loadWithNamedParam(h: h);
    }
    return result;
  }
}

class BImpl<T extends A> extends B<T> {
  BImpl(T v) : super(v);

  // These methods do not matter. They can just return 0.
  int load(H<T> h) => 0;
  int loadWithNamedParam({required H<T> h}) => 0;
}

class B0 extends B<A0> {
  B0(A0 a) : super(a);

  int load(H<A0> h) {
    return h.data.str.length;
  }

  int loadWithNamedParam({String? a, required H<A0> h, String? z}) {
    return h.data.str.length;
  }
}

class B1 extends B<A1> {
  B1(A1 a) : super(a);

  int load(H<A1> h) {
    return h.data.str.length;
  }

  int loadWithNamedParam({String? a, required H<A1> h, String? z}) {
    return h.data.str.length;
  }
}

@pragma('vm:never-inline')
@pragma('vm:testing:print-flow-graph', '*LICM')
int testNarrowingThroughIsCheckOnSubclass(B b) {
  int sum = 0;
  b.v.toString();
  for (var i = 0; i < 2; i++) {
    if (b is B1) {
      sum += b.v.str.length;
    }
  }
  return sum;
}

@pragma('vm:never-inline')
@pragma('vm:testing:print-flow-graph', '*LICM')
int testNarrowingThroughIsCheckWithTypeArgPolymorphic(B b) {
  int sum = 0;
  final v = b.v;
  for (var i = 0; i < 2; i++) {
    if (b is B<A1>) {
      sum += b.v.str.length;
    }
  }
  blackhole(v);
  return sum;
}

@pragma('vm:never-inline')
@pragma('vm:testing:print-flow-graph', '*LICM')
int testNarrowingThroughIsCheckWithTypeArgMonomorphic(B b) {
  int sum = 0;
  final v = b.v;
  for (var i = 0; i < 2; i++) {
    if (b is B<A1>) {
      sum += b.v.str.length;
    }
  }
  blackhole(v);
  return sum;
}

@pragma('vm:never-inline')
@pragma('vm:testing:print-flow-graph', '*LICM')
int testNarrowingThroughAsCheckWithTypeArgPolymorphic(B b) {
  int sum = 0;
  final v = b.v;
  for (var i = 0; i < 2; i++) {
    // Branch with a side-effect to avoid b as B<A1> hoisting.
    if (sum == 42) throw '42';
    sum += (b as B<A1>).v.str.length;
  }
  blackhole(v);
  return sum;
}

@pragma('vm:never-inline')
@pragma('vm:testing:print-flow-graph', '*LICM')
int testNarrowingThroughPhiOfAsChecksWithTypeArgPolymorphic(B b) {
  int sum = 0;
  final v = b.v;
  for (var i = 0; i < 2; i++) {
    // Branch with a side-effect to avoid b as B<A1> hoisting.
    if (sum == 42) throw '42';
    sum += (sum == 22 ? b as B<A1> : b as B<A1>).v.str.length;
  }
  blackhole(v);
  return sum;
}

@pragma('vm:never-inline')
@pragma('vm:testing:print-flow-graph', '*LICM')
int testNarrowingThroughIndexedLoadFromGrowableArray(List<A> l) {
  int sum = 0;
  final v = l[0];
  for (var i = 0; i < 2; i++) {
    if (l is List<A1>) {
      sum += l[0].str.length;
    }
  }
  blackhole(v);
  return sum;
}

@pragma('vm:never-inline')
@pragma('vm:testing:print-flow-graph', '*LICM')
int testNarrowingThroughIndexedLoadFromFixedArray(List<A> l) {
  int sum = 0;
  final v = l[0];
  for (var i = 0; i < 2; i++) {
    if (l is List<A1>) {
      sum += l[0].str.length;
    }
  }
  blackhole(v);
  return sum;
}

void main(List<String> args) {
  shouldPrint = args.contains("shouldPrint");

  // Prevent shaking of BImpl.load and BImpl.loadWithNamed, if these methods
  // are shaked (because they are not used) that would inhibit polymorphic
  // inlining at testNarrowingThroughThisCall{,WithNamedParams}
  BImpl(A1("")).load(H(A1("")));
  BImpl(A1("")).loadWithNamedParam(h: H(A1("")));

  for (var i = 0; i < 2; i++) {
    final a1 = A1("$i");
    final a0 = A0(i.toDouble(), "$i");
    B1(a1).testNarrowingThroughThisCallWithPositionalParam(H(a1));
    B0(a0).testNarrowingThroughThisCallWithPositionalParam(H(a0));
    B1(a1).testNarrowingThroughThisCallWithNamedParams(H(a1));
    B0(a0).testNarrowingThroughThisCallWithNamedParams(H(a0));
    testNarrowingThroughIsCheckOnSubclass(B1(a1));
    testNarrowingThroughIsCheckOnSubclass(B0(a0));
    testNarrowingThroughIsCheckWithTypeArgPolymorphic(B1(a1));
    testNarrowingThroughIsCheckWithTypeArgPolymorphic(B0(a0));
    testNarrowingThroughIsCheckWithTypeArgMonomorphic(BImpl<A1>(a1));
    testNarrowingThroughIsCheckWithTypeArgMonomorphic(BImpl<A0>(a0));
    testNarrowingThroughAsCheckWithTypeArgPolymorphic(B1(a1));
    try {
      testNarrowingThroughAsCheckWithTypeArgPolymorphic(B0(a0));
      throw "Should be unreachable";
    } on TypeError catch (e) {}
    testNarrowingThroughPhiOfAsChecksWithTypeArgPolymorphic(B1(a1));
    try {
      testNarrowingThroughPhiOfAsChecksWithTypeArgPolymorphic(B0(a0));
      throw "Should be unreachable";
    } on TypeError catch (e) {}
    testNarrowingThroughIndexedLoadFromGrowableArray([a1]);
    testNarrowingThroughIndexedLoadFromGrowableArray([a0]);
    testNarrowingThroughIndexedLoadFromFixedArray(
        List<A1>.filled(1, a1, growable: false));
    testNarrowingThroughIndexedLoadFromFixedArray(
        List<A0>.filled(1, a0, growable: false));
  }
}

void matchIL$B$testNarrowingThroughThisCallWithPositionalParam(
    FlowGraph beforeLICM, FlowGraph afterLICM) {
  final env = beforeLICM.match([
    match.block('Graph'),
    match.block('Function', [
      'this' << match.Parameter(index: 0),
      'h_raw' << match.Parameter(index: 1),
      'h' << match.AssertAssignable('h_raw', match.any),
      match.Goto('B1'),
    ]),
    'B1' <<
        match.block('Join', [
          match.CheckStackOverflow(),
          match.Branch(match.RelationalOp(match.any, match.any, kind: '<'),
              ifTrue: 'B2'),
        ]),
    'B2' <<
        match.block('Target', [
          'this_cid' << match.LoadClassId('this'),
          match.Branch(match.EqualityCompare('this_cid', match.any, kind: '=='),
              ifTrue: 'B3'),
        ]),
    'B3' <<
        match.block('Target', [
          // This redefinition was inserted by inlining.
          'h_' << match.Redefinition('h'),
          match.Redefinition('this'),
          'v0' << match.LoadField('h_', slot: 'data'),
          'v1' << match.LoadField('v0', slot: 'str'),
          'v2' << match.LoadField('v1', slot: 'String.length'),
        ]),
  ]);

  afterLICM.match([
    match.block('Graph'),
    match.block('Function', [
      'this' << match.Parameter(index: 0),
      'h_raw' << match.Parameter(index: 1),
      'h' << match.AssertAssignable('h_raw', match.any),
      'this_cid' << match.LoadClassId('this'), // Hoisted.
      match.Goto('B1'),
    ]),
    'B1' <<
        match.block('Join', [
          match.CheckStackOverflow(),
          match.Branch(match.RelationalOp(match.any, match.any, kind: '<'),
              ifTrue: 'B2'),
        ]),
    'B2' <<
        match.block('Target', [
          match.Branch(match.EqualityCompare('this_cid', match.any, kind: '=='),
              ifTrue: 'B3'),
        ]),
    'B3' <<
        match.block('Target', [
          // After LICM redefinitions are removed.
          'v0' << match.LoadField('h', slot: 'data'),
          'v1' << match.LoadField('v0', slot: 'str'),
          'v2' << match.LoadField('v1', slot: 'String.length'),
        ]),
  ], env: env);
}

void matchIL$B$testNarrowingThroughThisCallWithNamedParams(
    FlowGraph beforeLICM, FlowGraph afterLICM) {
  // Graph shape is basically the same.
  matchIL$B$testNarrowingThroughThisCallWithPositionalParam(
      beforeLICM, afterLICM);
}

void matchIL$testNarrowingThroughIsCheckOnSubclass(
    FlowGraph beforeLICM, FlowGraph afterLICM) {
  final env = beforeLICM.match([
    match.block('Graph'),
    match.block('Function', [
      'b' << match.Parameter(index: 0),
      'b.v' << match.LoadField('b', slot: 'v'),
      match.Goto('B1'),
    ]),
    'B1' <<
        match.block('Join', [
          match.CheckStackOverflow(),
          match.Branch(match.RelationalOp(match.any, match.any, kind: '<'),
              ifTrue: 'B2'),
        ]),
    'B2' <<
        match.block('Target', [
          'b_cid' << match.LoadClassId('b'),
          match.Branch(match.EqualityCompare('b_cid', match.any, kind: '=='),
              ifTrue: 'B3'),
        ]),
    'B3' <<
        match.block('Target', [
          match.Redefinition('b'),
          // This redefinition was inserted by load forwarding.
          'b.v*' << match.Redefinition('b.v'),
          'v0' << match.LoadField('b.v*', slot: 'str'),
          'v1' << match.LoadField('v0', slot: 'String.length'),
        ]),
  ]);

  afterLICM.match([
    match.block('Graph'),
    match.block('Function', [
      'b' << match.Parameter(index: 0),
      'b.v' << match.LoadField('b', slot: 'v'),
      'b_cid' << match.LoadClassId('b'), // Hoisted.
      match.Goto('B1'),
    ]),
    'B1' <<
        match.block('Join', [
          match.CheckStackOverflow(),
          match.Branch(match.RelationalOp(match.any, match.any, kind: '<'),
              ifTrue: 'B2'),
        ]),
    'B2' <<
        match.block('Target', [
          match.Branch(match.EqualityCompare('b_cid', match.any, kind: '=='),
              ifTrue: 'B3'),
        ]),
    'B3' <<
        match.block('Target', [
          'v0' << match.LoadField('b.v', slot: 'str'),
          'v1' << match.LoadField('v0', slot: 'String.length'),
        ]),
  ], env: env);
}

void matchIL$testNarrowingThroughIsCheckWithTypeArgMonomorphic(
    FlowGraph beforeLICM, FlowGraph afterLICM) {
  final env = beforeLICM.match([
    match.block('Graph'),
    match.block('Function', [
      'b' << match.Parameter(index: 0),
      'b.v' << match.LoadField('b', slot: 'v'),
      match.Goto('B1'),
    ]),
    'B1' <<
        match.block('Join', [
          match.CheckStackOverflow(),
          match.Branch(match.RelationalOp(match.any, match.any, kind: '<'),
              ifTrue: 'B2'),
        ]),
    'B2' <<
        match.block('Target', [
          'b_is_B<A1>' << match.InstanceOf('b', match.any, match.any),
          match.Branch(
              match.StrictCompare('b_is_B<A1>', match.any, kind: '==='),
              ifTrue: 'B3'),
        ]),
    'B3' <<
        match.block('Target', [
          match.Redefinition('b'),
          // This redefinition was inserted by load forwarding.
          'b.v*' << match.Redefinition('b.v'),
          'v0' << match.LoadField('b.v*', slot: 'str'),
          'v1' << match.LoadField('v0', slot: 'String.length'),
        ]),
  ]);

  afterLICM.match([
    match.block('Graph'),
    match.block('Function', [
      'b' << match.Parameter(index: 0),
      'b.v' << match.LoadField('b', slot: 'v'),
      match.Goto('B1'),
    ]),
    'B1' <<
        match.block('Join', [
          match.CheckStackOverflow(),
          match.Branch(match.RelationalOp(match.any, match.any, kind: '<'),
              ifTrue: 'B2'),
        ]),
    'B2' <<
        match.block('Target', [
          'b_is_B<A1>' << match.InstanceOf('b', match.any, match.any),
          match.Branch(
              match.StrictCompare('b_is_B<A1>', match.any, kind: '==='),
              ifTrue: 'B3'),
        ]),
    'B3' <<
        match.block('Target', [
          'v0' << match.LoadField('b.v', slot: 'str'),
          'v1' << match.LoadField('v0', slot: 'String.length'),
        ]),
  ], env: env);
}

void matchIL$testNarrowingThroughAsCheckWithTypeArgPolymorphic(
    FlowGraph beforeLICM, FlowGraph afterLICM) {
  final env = beforeLICM.match([
    match.block('Graph'),
    match.block('Function', [
      'b' << match.Parameter(index: 0),
      'b.v' << match.LoadField('b', slot: 'v'),
      match.Goto('B1'),
    ]),
    'B1' <<
        match.block('Join', [
          match.CheckStackOverflow(),
          match.Branch(match.RelationalOp(match.any, match.any, kind: '<'),
              ifTrue: 'B2'),
        ]),
    'B2' <<
        match.block('Target', [
          match.Branch(match.EqualityCompare(match.any, match.any, kind: '=='),
              ifTrue: 'B3', ifFalse: 'B4'),
        ]),
    'B3' << match.block('Target', [match.Throw(match.any)]),
    'B4' <<
        match.block('Target', [
          match.AssertAssignable('b', match.any, match.any, match.any),
          // This redefinition was inserted by load forwarding.
          'b.v*' << match.Redefinition('b.v'),
          'v0' << match.LoadField('b.v*', slot: 'str'),
          'v1' << match.LoadField('v0', slot: 'String.length'),
        ]),
  ]);

  afterLICM.match([
    match.block('Graph'),
    match.block('Function', [
      'b' << match.Parameter(index: 0),
      'b.v' << match.LoadField('b', slot: 'v'),
      match.Goto('B1'),
    ]),
    'B1' <<
        match.block('Join', [
          match.CheckStackOverflow(),
          match.Branch(match.RelationalOp(match.any, match.any, kind: '<'),
              ifTrue: 'B2'),
        ]),
    'B2' <<
        match.block('Target', [
          match.Branch(match.EqualityCompare(match.any, match.any, kind: '=='),
              ifTrue: 'B3', ifFalse: 'B4'),
        ]),
    'B3' << match.block('Target', [match.Throw(match.any)]),
    'B4' <<
        match.block('Target', [
          match.AssertAssignable('b', match.any, match.any, match.any),
          'v0' << match.LoadField('b.v', slot: 'str'),
          'v1' << match.LoadField('v0', slot: 'String.length'),
        ]),
  ], env: env);
}

void matchIL$testNarrowingThroughPhiOfAsChecksWithTypeArgPolymorphic(
    FlowGraph beforeLICM, FlowGraph afterLICM) {
  final env = beforeLICM.match([
    match.block('Graph'),
    match.block('Function', [
      'b' << match.Parameter(index: 0),
      'b.v' << match.LoadField('b', slot: 'v'),
      match.Goto('B1'),
    ]),
    'B1' <<
        match.block('Join', [
          match.CheckStackOverflow(),
          match.Branch(match.RelationalOp(match.any, match.any, kind: '<'),
              ifTrue: 'B2'),
        ]),
    'B2' <<
        match.block('Target', [
          match.Branch(match.EqualityCompare(match.any, match.any, kind: '=='),
              ifTrue: 'B3', ifFalse: 'B4'),
        ]),
    'B3' << match.block('Target', [match.Throw(match.any)]),
    'B4' <<
        match.block('Target', [
          match.Branch(match.EqualityCompare(match.any, match.any, kind: '=='),
              ifTrue: 'B5', ifFalse: 'B6'),
        ]),
    'B5' <<
        match.block('Target', [
          match.AssertAssignable('b', match.any, match.any, match.any),
          match.Goto('B7'),
        ]),
    'B6' <<
        match.block('Target', [
          match.AssertAssignable('b', match.any, match.any, match.any),
          match.Goto('B7'),
        ]),
    'B7' <<
        match.block('Join', [
          // This redefinition was inserted by PhiInstr::Canonicalize
          // which removed Phi of two AssertAssignables above.
          match.Redefinition('b'),
          // This redefinition was inserted by load forwarding.
          'b.v*' << match.Redefinition('b.v'),
          'v0' << match.LoadField('b.v*', slot: 'str'),
          'v1' << match.LoadField('v0', slot: 'String.length'),
        ]),
  ]);

  afterLICM.match([
    match.block('Graph'),
    match.block('Function', [
      'b' << match.Parameter(index: 0),
      'b.v' << match.LoadField('b', slot: 'v'),
      match.Goto('B1'),
    ]),
    'B1' <<
        match.block('Join', [
          match.CheckStackOverflow(),
          match.Branch(match.RelationalOp(match.any, match.any, kind: '<'),
              ifTrue: 'B2'),
        ]),
    'B2' <<
        match.block('Target', [
          match.Branch(match.EqualityCompare(match.any, match.any, kind: '=='),
              ifTrue: 'B3', ifFalse: 'B4'),
        ]),
    'B3' << match.block('Target', [match.Throw(match.any)]),
    'B4' <<
        match.block('Target', [
          match.Branch(match.EqualityCompare(match.any, match.any, kind: '=='),
              ifTrue: 'B5', ifFalse: 'B6'),
        ]),
    'B5' <<
        match.block('Target', [
          match.AssertAssignable('b', match.any, match.any, match.any),
          match.Goto('B7'),
        ]),
    'B6' <<
        match.block('Target', [
          match.AssertAssignable('b', match.any, match.any, match.any),
          match.Goto('B7'),
        ]),
    'B7' <<
        match.block('Join', [
          'v0' << match.LoadField('b.v', slot: 'str'),
          'v1' << match.LoadField('v0', slot: 'String.length'),
        ]),
  ], env: env);
}

void matchIL$testNarrowingThroughIndexedLoadFromGrowableArray(
    FlowGraph beforeLICM, FlowGraph afterLICM) {
  final env = beforeLICM.match([
    match.block('Graph'),
    match.block('Function', [
      'this' << match.Parameter(index: 0),
      'a.data' << match.LoadField('this', slot: 'GrowableObjectArray.data'),
      'a.data[0]' << match.LoadIndexed('a.data', match.any),
      match.Goto('B1'),
    ]),
    'B1' <<
        match.block('Join', [
          match.CheckStackOverflow(),
          match.Branch(match.RelationalOp(match.any, match.any, kind: '<'),
              ifTrue: 'B2'),
        ]),
    'B2' <<
        match.block('Target', [
          'this_is_List' << match.InstanceOf('this', match.any, match.any),
          match.Branch(
              match.StrictCompare('this_is_List', match.any, kind: '==='),
              ifTrue: 'B3'),
        ]),
    'B3' <<
        match.block('Target', [
          match.Redefinition('this'),
          // This redefinition was inserted by load forwarding.
          'a.data[0]*' << match.Redefinition('a.data[0]'),
          'a.data[0].str' << match.LoadField('a.data[0]*', slot: 'str'),
        ]),
  ]);

  afterLICM.match([
    match.block('Graph'),
    match.block('Function', [
      'this' << match.Parameter(index: 0),
      'a.data' << match.LoadField('this', slot: 'GrowableObjectArray.data'),
      'a.data[0]' << match.LoadIndexed('a.data', match.any),
      match.Goto('B1'),
    ]),
    'B1' <<
        match.block('Join', [
          match.CheckStackOverflow(),
          match.Branch(match.RelationalOp(match.any, match.any, kind: '<'),
              ifTrue: 'B2'),
        ]),
    'B2' <<
        match.block('Target', [
          'this_is_List' << match.InstanceOf('this', match.any, match.any),
          match.Branch(
              match.StrictCompare('this_is_List', match.any, kind: '==='),
              ifTrue: 'B3'),
        ]),
    'B3' <<
        match.block('Target', [
          'a.data[0].str' << match.LoadField('a.data[0]', slot: 'str'),
        ]),
  ], env: env);
}

void matchIL$testNarrowingThroughIsCheckWithTypeArgPolymorphic(
    FlowGraph beforeLICM, FlowGraph afterLICM) {
  matchIL$testNarrowingThroughIsCheckWithTypeArgMonomorphic(
      beforeLICM, afterLICM);
}

void matchIL$testNarrowingThroughIndexedLoadFromFixedArray(
    FlowGraph beforeLICM, FlowGraph afterLICM) {
  final env = beforeLICM.match([
    match.block('Graph'),
    match.block('Function', [
      'this' << match.Parameter(index: 0),
      'a[0]' << match.LoadIndexed('this', match.any),
      match.Goto('B1'),
    ]),
    'B1' <<
        match.block('Join', [
          match.CheckStackOverflow(),
          match.Branch(match.RelationalOp(match.any, match.any, kind: '<'),
              ifTrue: 'B2'),
        ]),
    'B2' <<
        match.block('Target', [
          'this_is_List' << match.InstanceOf('this', match.any, match.any),
          match.Branch(
              match.StrictCompare('this_is_List', match.any, kind: '==='),
              ifTrue: 'B3'),
        ]),
    'B3' <<
        match.block('Target', [
          match.Redefinition('this'),
          // This redefinition was inserted by load forwarding.
          'a[0]*' << match.Redefinition('a[0]'),
          'a[0].str' << match.LoadField('a[0]*', slot: 'str'),
        ]),
  ]);

  afterLICM.match([
    match.block('Graph'),
    match.block('Function', [
      'this' << match.Parameter(index: 0),
      'a[0]' << match.LoadIndexed('this', match.any),
      match.Goto('B1'),
    ]),
    'B1' <<
        match.block('Join', [
          match.CheckStackOverflow(),
          match.Branch(match.RelationalOp(match.any, match.any, kind: '<'),
              ifTrue: 'B2'),
        ]),
    'B2' <<
        match.block('Target', [
          'this_is_List' << match.InstanceOf('this', match.any, match.any),
          match.Branch(
              match.StrictCompare('this_is_List', match.any, kind: '==='),
              ifTrue: 'B3'),
        ]),
    'B3' <<
        match.block('Target', [
          'a[0].str' << match.LoadField('a[0]', slot: 'str'),
        ]),
  ], env: env);
}
